package org.example.graph;

import java.util.*;

public class Prim {




    public static class EdgeComparator implements Comparator<Edge> {


        @Override
        public int compare(Edge o1, Edge o2) {
            return o1.weight - o2.weight;
        }
    }

    public static Set<Edge> PrimMST(Graph graph) {
        PriorityQueue<Edge> priorityQueue = new PriorityQueue<>(new EdgeComparator());

        HashSet<Node> set = new HashSet<>();

        Set<Edge> result = new HashSet<>();
        for (Node node : graph.nodes.values()) {

            // node  开始的点
            if (!set.contains(node)) {
                set.add(node);
                for (Edge edge : node.edges) {
                    // node 所有的边放入队列里面
                    priorityQueue.add(edge);
                }
                while (!priorityQueue.isEmpty()) {
                    Edge edge = priorityQueue.poll(); //弹出最小边
                    Node toNode = edge.to;  // 是否为一个新的点
                    if (!set.contains(toNode)) {
                        set.add(toNode);
                        result.add(edge);
                        for (Edge nextEdge : toNode.edges) {
                            priorityQueue.add(nextEdge);
                        }
                    }
                }

            }

        }
        return result;
    }



}
